/* 
 * Copyright 2014 Josef Hertl.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package caetus.pkcs11.template;

import caetus.pkcs11.key.CipherEnum;
import java.math.BigInteger;
import java.security.NoSuchAlgorithmException;
import java.security.spec.InvalidKeySpecException;
import java.util.Random;
import sun.security.pkcs11.wrapper.CK_ATTRIBUTE;
import sun.security.pkcs11.wrapper.PKCS11Constants;

/**
 *
 * @author Josef Hertl
 */
public class TemplateMaker {

    private Random random = new Random();

    /**
     * Generates an array of random bytes.
     *
     * @param length The length of the generated bytes[] array
     * @return A newly generated array of random bytes.
     */
    private byte[] generateBytes(int length) {
        byte[] generatedBytes = new byte[length];
        random.nextBytes(generatedBytes);
        return generatedBytes;
    }

    /**
     * Generates the template for the given key type, with the required boolean
     * values. Key value is randomly generated based off the key type. Does NOT
     * include any generated value of the key.
     *
     * @param keyType The algorithm for which the secret key will work.
     * @param sensitive Whether the key value should never be revealed.
     * @param extractable Whether the can be extracted (=exported) from the
     * token
     * @param encrypt If the key can be used for encrypting.
     * @param decrypt If the key can be used for decryption.
     * @param wrap Whether the key can be used to wrap other keys.
     * @param unwrap Whether the key can be used to unwrap other keys.
     * @return The new secret key template based on the above values.
     */
    public Template getSecretKeyTemplate(CipherEnum cipher, boolean sensitive, boolean extractable, boolean encrypt, boolean decrypt, boolean wrap, boolean unwrap) {
        Template template = new Template();
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_SENSITIVE, sensitive));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_EXTRACTABLE, extractable));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_CLASS, PKCS11Constants.CKO_SECRET_KEY));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_KEY_TYPE, cipher.getKeyType()));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_ENCRYPT, encrypt));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_DECRYPT, decrypt));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_WRAP, wrap));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_UNWRAP, unwrap));
        return template;
    }

    /**
     * @return Basic template with all the values set to true.
     */
    public Template getBasicSecretKeyTemplate(CipherEnum cipherEnum) {
        Template template = new Template();
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_SENSITIVE, true));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_EXTRACTABLE, true));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_CLASS, PKCS11Constants.CKO_SECRET_KEY));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_KEY_TYPE, cipherEnum.getKeyType()));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_ENCRYPT, true));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_DECRYPT, true));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_WRAP, true));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_UNWRAP, true));
        return template;
    }

    /**
     * @return Template that marks the key as sensitive and extractable and only
     * permits encryption and decryption.
     */
    public Template getSecretKeyEncryptDecryptTemplate(CipherEnum cipherEnum) {
        Template template = new Template();
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_SENSITIVE, true));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_EXTRACTABLE, true));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_CLASS, PKCS11Constants.CKO_SECRET_KEY));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_KEY_TYPE, cipherEnum.getKeyType()));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_ENCRYPT, true));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_DECRYPT, true));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_WRAP, false));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_UNWRAP, false));
        return template;
    }

    /**
     * @return Template that only permits wrap and marks the key as extractable
     * and sensitive.
     */
    public Template getSecretKeyWrapUnwrapTemplate(CipherEnum cipherEnum) {
        Template template = new Template();
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_SENSITIVE, true));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_EXTRACTABLE, true));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_CLASS, PKCS11Constants.CKO_SECRET_KEY));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_KEY_TYPE, cipherEnum.getKeyType()));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_ENCRYPT, false));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_DECRYPT, false));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_WRAP, true));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_UNWRAP, true));
        return template;
    }

    /**
     *
     * @return Public key template for generating RSA keys.
     * @param modulusBits Recommended to set to 3072.
     * @throws NoSuchAlgorithmException
     * @throws InvalidKeySpecException
     */
    public Template getRSAPublicKeyGenerationTemplate(byte[] pubExponent, int modulusBits, boolean encrypt, boolean wrap) throws NoSuchAlgorithmException, InvalidKeySpecException {
        Template template = new Template();
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_CLASS, PKCS11Constants.CKO_PUBLIC_KEY));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_KEY_TYPE, PKCS11Constants.CKK_RSA));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_ENCRYPT, encrypt));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_WRAP, wrap));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_MODULUS_BITS, modulusBits));
//        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_PUBLIC_EXPONENT, new byte[cipher.getPublicValue().length]));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_PUBLIC_EXPONENT, pubExponent));
        return template;
    }

    /**
     * @return Private key template for generating RSA keys, but without
     * MODULUS_BITS.
     * @throws NoSuchAlgorithmException
     * @throws InvalidKeySpecException
     */
    public Template getRSAPrivateKeyGenerationTemplate(boolean sensitive, boolean extractable, boolean decrypt, boolean unwrap) throws NoSuchAlgorithmException, InvalidKeySpecException {
        Template template = new Template();
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_CLASS, PKCS11Constants.CKO_PRIVATE_KEY));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_KEY_TYPE, PKCS11Constants.CKK_RSA));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_DECRYPT, decrypt));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_UNWRAP, unwrap));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_SENSITIVE, sensitive));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_EXTRACTABLE, extractable));
        return template;
    }

    /**
     * @return A RSA template with all values set to extremely weak values.
     */
    public Template getBasicRSAPublicKeyTemplate(boolean encrypt, boolean wrap) throws NoSuchAlgorithmException, InvalidKeySpecException {
        Template template = new Template();
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_CLASS, PKCS11Constants.CKO_PUBLIC_KEY));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_KEY_TYPE, PKCS11Constants.CKK_RSA));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_ENCRYPT, encrypt));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_WRAP, wrap));
//        BigInteger bigModulus = new BigInteger("2501");
        BigInteger bigModulus = BigInteger.probablePrime(2000, new Random());
        BigInteger bigPublic = new BigInteger("23");
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_MODULUS, bigModulus.toByteArray()));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_PUBLIC_EXPONENT, bigPublic.toByteArray()));
        return template;
    }

    /**
     * @return RSA template that only contains preset and pregenerated public
     * key value.
     */
    public Template getBasicRSAPublicKeyValueTemplate() {
        Template template = new Template();
        BigInteger bigModulus = BigInteger.probablePrime(2000, new Random());
        BigInteger bigPublic = new BigInteger("23");
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_MODULUS, bigModulus.toByteArray()));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_PUBLIC_EXPONENT, bigPublic.toByteArray()));
        return template;
    }

    /**
     * @return Basic private key template with very weak settings.
     * @throws NoSuchAlgorithmException
     * @throws InvalidKeySpecException
     */
    public Template getBasicRSAPrivateKeyTemplate(boolean sensitive, boolean extractable, boolean decrypt, boolean unwrap) throws NoSuchAlgorithmException, InvalidKeySpecException {
        Template template = new Template();
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_CLASS, PKCS11Constants.CKO_PRIVATE_KEY));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_KEY_TYPE, PKCS11Constants.CKK_RSA));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_DECRYPT, decrypt));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_UNWRAP, unwrap));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_SENSITIVE, sensitive));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_EXTRACTABLE, extractable));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_LABEL, "rsaPriv"));
        BigInteger bigModulus = new BigInteger("2501");
        BigInteger bigPrivate = new BigInteger("2087");
        BigInteger bigPublic = new BigInteger("23");
        BigInteger bigPrime1 = new BigInteger("41");
        BigInteger bigPrime2 = new BigInteger("61");
        BigInteger bigExponent1 = new BigInteger("7");
        BigInteger bigExponent2 = new BigInteger("47");
        BigInteger bigCoefficient = new BigInteger("39");
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_MODULUS, bigModulus.toByteArray()));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_PUBLIC_EXPONENT, bigPublic.toByteArray()));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_PRIVATE_EXPONENT, bigPrivate.toByteArray()));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_PRIME_1, bigPrime1.toByteArray()));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_PRIME_2, bigPrime2.toByteArray()));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_EXPONENT_1, bigExponent1.toByteArray()));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_EXPONENT_2, bigExponent2.toByteArray()));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_COEFFICIENT, bigCoefficient.toByteArray()));
        return template;
    }

    /**
     * @return RSA private key template with only key value.
     */
    public Template getBasicRSAPrivateKeyValueTemplate() throws NoSuchAlgorithmException, InvalidKeySpecException {
        Template template = new Template();
        BigInteger bigModulus = new BigInteger("2501");
        BigInteger bigPrivate = new BigInteger("2087");
        BigInteger bigPublic = new BigInteger("23");
        BigInteger bigPrime1 = new BigInteger("41");
        BigInteger bigPrime2 = new BigInteger("61");
        BigInteger bigExponent1 = new BigInteger("7");
        BigInteger bigExponent2 = new BigInteger("47");
        BigInteger bigCoefficient = new BigInteger("39");
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_MODULUS, bigModulus.toByteArray()));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_PUBLIC_EXPONENT, bigPublic.toByteArray()));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_PRIVATE_EXPONENT, bigPrivate.toByteArray()));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_PRIME_1, bigPrime1.toByteArray()));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_PRIME_2, bigPrime2.toByteArray()));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_EXPONENT_1, bigExponent1.toByteArray()));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_EXPONENT_2, bigExponent2.toByteArray()));
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_COEFFICIENT, bigCoefficient.toByteArray()));
        return template;
    }

    /**
     * @return Template that only has one feature - value. It is useful for
     * finding values of various objects.
     */
    public Template getSimpleValueTemplate() {
        Template template = new Template();
        template.addAttribute(new CK_ATTRIBUTE(PKCS11Constants.CKA_VALUE));
        return template;
    }
}
